import json
import logging
import os
import sys
import time

from _dynamic_dataset.TrainDataset import TrainSubset
from rootPath import project_path
from server_client.Client import Client
from server_client.Server import Server
import torch
from torch.utils.data import DataLoader
import torchvision
import torchvision.transforms as transforms
from datetime import datetime

from utils.datasetUtils import splitTrainDataByNums
from utils.modelUtils import getResNet18, getAvailableDevice

'''
实验目的：验证各方数据集在相同大小和相同分布的情况下的真实夏普利值的计算
'''

# 获取当前日期和时间
now = datetime.now()
# 超参数
batch_size = 128
learning_rate = 0.0001
num_epochs = 200
# threshold = args.threshold
# maxNum = args.maxNum

# 读取配置文件
# 基于梯度的夏普利值计算
with open(project_path + "/conf/client.json", 'r') as f:
    clientConf = json.load(f)

with open(project_path + "/conf/server.json", 'r') as f:
    serverConf = json.load(f)

# 数据预处理
train_transform = transforms.Compose(
    [
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))])

# 测试集数据预处理
test_transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
])

# 客户端的数据分配创建
numsList = [
    {0: 1000, 1: 1000, 2: 1000, 3: 1000, 4: 1000, 5: 1000, 6: 1000, 7: 1000, 8: 1000, 9: 1000},
    {0: 1000, 1: 1000, 2: 1000, 3: 1000, 4: 1000, 5: 1000, 6: 1000, 7: 1000, 8: 1000, 9: 1000},
    {0: 1000, 1: 1000, 2: 1000, 3: 1000, 4: 1000, 5: 1000, 6: 1000, 7: 1000, 8: 1000, 9: 1000},
    {0: 1000, 1: 1000, 2: 1000, 3: 1000, 4: 1000, 5: 1000, 6: 1000, 7: 1000, 8: 1000, 9: 1000},
    {0: 1000, 1: 1000, 2: 1000, 3: 1000, 4: 1000, 5: 1000, 6: 1000, 7: 1000, 8: 1000, 9: 1000},
]
clientNum = len(numsList)
clients = []
print("生成id组合完成")

device = getAvailableDevice()
# 服务器的验证集
test_set = torchvision.datasets.CIFAR10(root=project_path + '/data', train=False, download=True,
                                        transform=test_transform)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=2)
# 测试集
train_dataset = torchvision.datasets.CIFAR10(root=project_path + '/data/', train=True, download=True,
                                             transform=train_transform)
model = getResNet18()
datasetIndices = splitTrainDataByNums(orderList=numsList, trainDataset=train_dataset)
for i in range(clientNum):
    trainSubDataset = TrainSubset(train_dataset, datasetIndices[0])
    clients.append(
        Client(clientConf, model, device, DataLoader(trainSubDataset, shuffle=True, batch_size=128, num_workers=2), i))
    print(f"客户端{i}创建成功")
print("客户端全部初始化成功！")

# 服务器
server = Server(serverConf, test_loader, device, model)

for epoch in range(1, num_epochs + 1):
    logging.info(f"epoch={epoch}")
    # 客户端训练
    for client in clients:
        client.train(server.global_model)
    # 服务器梯度聚合并计算对应的正确率
    server.model_aggregate(clients, server.global_model)
    acc, loss = server.model_eval(server.global_model, server.eval_loader)
    print(f"epoch={epoch},acc={acc},loss={loss}")
